package edu.zju.gis.dbfg.model.tile;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.common.util.DateUtils;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.base.GeoItem;
import edu.zju.gis.dbfg.model.etl.OptionField;
import edu.zju.gis.dbfg.model.scala.MIngest;
import edu.zju.gis.dbfg.model.scala.MRasterizer;
import edu.zju.gis.dbfg.model.statistic.AreaStatistic;
import edu.zju.gis.dbfg.model.util.FormatUtil;
import geotrellis.raster.MultibandTile;
import geotrellis.raster.RasterExtent;
import geotrellis.raster.render.ColorMap;
import geotrellis.vector.Extent;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.*;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.storage.StorageLevel;
import org.elasticsearch.spark.rdd.api.java.JavaEsSpark;
import org.geotools.geojson.geom.GeometryJSON;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.math.BigInteger;
import java.nio.file.FileAlreadyExistsException;
import java.sql.*;
import java.util.*;

public class TileServiceFromEs {
    private final static Logger logger = LoggerFactory.getLogger(TileServiceFromEs.class);
    private static final String WGS84_WKT = "GEOGCS[\"GCS_WGS_1984\", \n" +
            "  DATUM[\"D_WGS_1984\", \n" +
            "    SPHEROID[\"WGS_1984\", 6378137.0, 298.257223563]], \n" +
            "  PRIMEM[\"Greenwich\", 0.0], \n" +
            "  UNIT[\"degree\", 0.017453292519943295], \n" +
            "  AXIS[\"Longitude\", EAST], \n" +
            "  AXIS[\"Latitude\", NORTH]]";
    private static String publisher = "admin";
    private static String tileServerTail = "/geotrellis/{z}/{x}/{y}/";//with no band
    private static Integer defaultZNumb = 22;
    private static String userSpaceRoot = "/opt/shengting/userSpace/";

    public static void main(String[] args) throws Exception {

        // Setup model arguments
        logger.info("Setup arguments");
        TileServiceFromEsArgs mArgs = TileServiceFromEsArgs.initArgs(args);
        if (mArgs == null) {
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(AreaStatistic.class, "main()", msg, args);
        }

        String id = UUID.randomUUID().toString();
        String indexName = mArgs.getIndexName();//input
        String tileServerAddress = mArgs.getTileServerAddress();//input
        String layerName = mArgs.getLayerName();//input
        String fieldsAarry = mArgs.getFieldsAarry();//input
        String keyWords = mArgs.getKeyWords();//input
        String topic = mArgs.getTopic();//input
        String serviceURL = tileServerAddress + id + tileServerTail;
        String esIndexName = indexName;
        String output = userSpaceRoot + mArgs.getUserId() + "/mapServices";
        String outputPath = output + File.separator + layerName;
//        String output = "C://Data//dbfg//test0423";
//        String outputPath = "C://Data//dbfg//test0423//" + layerName;

        File outputFile = new File(outputPath);
        if (outputFile.exists()) {
            logger.error("Output file has already exists: " + outputPath);
            throw new FileAlreadyExistsException(outputPath);
        }

        String[] fields = fieldsAarry.split("#");
        OptionField[] feaFields = new OptionField[fields.length];
        for (int i = 0; i < feaFields.length; i++) {
            feaFields[i] = new OptionField(fields[i]);
        }

        StringBuffer field_info = new StringBuffer();
        field_info.append("{\"fields\":[");
        for (int i = 0; i < feaFields.length; i++) {
            if (i != 0)
                field_info.append(",");
            field_info.append("{\"name\":\"" + feaFields[i].getFieldName() + "\",\"type\":\"" + feaFields[i].getFieldType() + "\"}");
        }
        field_info.append("]}");

        Connection conn;
        PreparedStatement stmt;
        String driver = "com.mysql.cj.jdbc.Driver";
        String url = "jdbc:mysql://101.132.255.121:3306/shengting?allowMultiQueries=true&useUnicode=true&characterEncoding=UTF-8&useSSL=false&serverTimezone=GMT%2B8";
        String user = "root";
        String password = "root";
        String sql = "insert into tb_service values (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)";
        String sqlfield = "insert into tb_service_field_info values (?,?,?,?,?,?,?)";
        String updateSql = "update tb_service set status = ? where id = ?";

        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf()
//                .setMaster("local[1]")
//                .setAppName("Map Service Example")
                .set("es.index.auto.create", "true")
                .set("es.index.read.missing.as.empty", "true")
                .set("es.nodes.wan.only", "true")
                .set("es.nodes", "121.89.211.142")
                .set("es.port", "9200");
        JavaSparkContext jsc = new JavaSparkContext(conf);
        PyramidConfig pc = new PyramidConfig.PyramidConfigBuilder()
                .setBaseMapEnv(-180, 180, -90, 90)
                .setCrs(CRS.parseWKT(WGS84_WKT))
                .setzLevelRange(4, 22)
                .build();

        Broadcast<PyramidConfig> pcBroad = jsc.broadcast(pc);
        Broadcast<Integer> zBroad = jsc.broadcast(defaultZNumb);

        /**
         * 从 ES 中读取数据并转化为GeoItem
         */
        JavaPairRDD<String, Map<String, Object>> features = JavaEsSpark.esRDD(jsc, esIndexName);
        JavaRDD<GeoItem> data = features.map(new Function<Tuple2<String, Map<String, Object>>, GeoItem>() {
            @Override
            public GeoItem call(Tuple2<String, Map<String, Object>> feature) throws Exception {

                GeoItem gi = new GeoItem();
                Object shpFromEs = feature._2.get("shape");
                String geoJson = FormatUtil.Shape2GeoJson(shpFromEs);
                GeometryJSON g = new GeometryJSON();
                Geometry geometry = g.read(geoJson);
                gi.setGeometry(geometry);
                Map<String, Object> props = new HashMap<>();
                for (int i = 0; i < feaFields.length; i++) {
                    props.put(feaFields[i].getFieldName(), feature._2.get(feaFields[i].getFieldName()));
                }
                gi.setProp(props);
                gi.setFid(feature._1);
                return gi;
            }
        });

        data.persist(StorageLevel.MEMORY_AND_DISK());

        /**
         * 对于 Field 类型为 DDouble (离散 Double) 或者 CDouble (连续 Double)，计算得到值的个数或值域范围
         * 对于 Field 类型为 DString (离散 String) 的字段，计算所有可能的值，并构建 String -> Int 的映射表
         */
        Map<OptionField, String> fieldValueInfo = data.flatMapToPair(new PairFlatMapFunction<GeoItem, OptionField, Tuple2<OptionField, String>>() {
            @Override
            public Iterator<Tuple2<OptionField, Tuple2<OptionField, String>>> call(GeoItem gi) throws Exception {
                List<Tuple2<OptionField, Tuple2<OptionField, String>>> result = new ArrayList<>();
                for (OptionField f : feaFields) {
                    String fname = f.getFieldName();
                    String fvalue = String.valueOf(gi.getProp().get(fname));
                    result.add(new Tuple2<>(f, new Tuple2<>(f, fvalue)));
                }
                return result.iterator();
            }
        }).reduceByKey(new Function2<Tuple2<OptionField, String>, Tuple2<OptionField, String>, Tuple2<OptionField, String>>() {
            @Override
            public Tuple2<OptionField, String> call(Tuple2<OptionField, String> in1, Tuple2<OptionField, String> in2) throws Exception {
                String result = null;
                if (in1._1.getFieldType().equals("DString") || in1._1.getFieldType().equals("DDouble")) {
                    List<String> v1s = Arrays.asList(in1._2.split("##"));
                    List<String> r = new ArrayList<>();
                    r.addAll(v1s);
                    String[] v2s = in2._2.split("##");
                    for (String v : v2s) {
                        if (!v1s.contains(v)) {
                            r.add(v);
                        }
                    }
                    if (r.size() > 500) {
                        r.clear();
                        r.add("EMPTY");
                    }
                    result = String.join("##", r);
                } else if (in1._1.getFieldType().equals("CDouble")) {
                    String[] d1s = in1._2.split("##");
                    String[] d2s = in2._2.split("##");
                    Double min1 = Double.valueOf(d1s[0]);
                    Double max1 = Double.valueOf(d1s[d1s.length - 1]);
                    Double min2 = Double.valueOf(d2s[0]);
                    Double max2 = Double.valueOf(d2s[d2s.length - 1]);
                    Double min = min1 < min2 ? min1 : min2;
                    Double max = max1 > max2 ? max1 : max2;
                    result = String.format("%f##%f", min, max);
                } else {
                    logger.error("Unvalid OptionField Type, only support: DString, DDouble, CDouble");
                }
                return new Tuple2<>(in1._1, result);
            }
        }).mapToPair(new PairFunction<Tuple2<OptionField, Tuple2<OptionField, String>>, OptionField, String>() {
            @Override
            public Tuple2<OptionField, String> call(Tuple2<OptionField, Tuple2<OptionField, String>> in) throws Exception {
                if (in._2._2.equals("EMPTY")) {
                    OptionField of = new OptionField();
                    of.setFieldName(in._1.getFieldName());
                    of.setFieldType("EMPTY");
                    return new Tuple2<>(of, null);
                } else {
                    return in._2;
                }
            }
        }).collectAsMap();

        Map<String, Map<String, Integer>> valueTrans = new HashMap<>();
        for (OptionField key : fieldValueInfo.keySet()) {
            if (!key.getFieldType().equals("DString")) continue;
            if (fieldValueInfo.get(key).equals("EMPTY")) {
                logger.error("Unvalid field for render: too many discrete value for field with value " + key);
                continue;
            }
            String[] v = fieldValueInfo.get(key).split("##");
            Map<String, Integer> ft = new HashMap<>();
            for (int i = 0; i < v.length; i++) {
                ft.put(v[i], i + 1);
            }
            valueTrans.put(key.getFieldName(), ft);
        }

        Broadcast<Map<String, Map<String, Integer>>> valueTransBroad = jsc.broadcast(valueTrans);
        Broadcast<Map<OptionField, String>> fieldValueInfoBroad = jsc.broadcast(fieldValueInfo);

        JavaRDD<Double[]> extent = data.map(new Function<GeoItem, GeoItem>() {
            @Override
            public GeoItem call(GeoItem gi) throws Exception {
                Map<String, Object> props = gi.getProp();
                for (OptionField field : feaFields) {
                    if (valueTrans.get(field.getFieldName()) == null) {
                        // 说明是字段值过多的字段或者不是DString类型字段, 不做处理
                        continue;
                    }
                    String ov = String.valueOf(props.get(field.getFieldName()));
                    if (ov != null) {
                        Integer fv = valueTransBroad.getValue().get(field.getFieldName()).get(ov);
                        props.put(field.getFieldName(), fv);
                    }
                }
                gi.setProp(props);
                return gi;
            }
        }).repartition(1).mapPartitions(new FlatMapFunction<Iterator<GeoItem>, Double[]>() {
            /**
             * 计算四至
             * @param gii
             * @return
             * @throws Exception
             */
            @Override
            public Iterator<Double[]> call(Iterator<GeoItem> gii) throws Exception {

                List<GeoItem> gis = new ArrayList<>();
                double xmin = 180d;
                double xmax = -180d;
                double ymin = 90d;
                double ymax = -90d;
                int count = 0;
                while (gii.hasNext()) {
                    count++;
                    GeoItem feature = gii.next();
                    Geometry geometry = feature.getGeometry();
                    Envelope e = geometry.getEnvelopeInternal();
                    xmin = e.getMinX() < xmin ? e.getMinX() : xmin;
                    xmax = e.getMaxX() > xmax ? e.getMaxX() : xmax;
                    ymin = e.getMinY() < ymin ? e.getMinY() : ymin;
                    ymax = e.getMaxY() > ymax ? e.getMaxY() : ymax;

                    gis.add(feature);
                }

                Extent extent = new Extent(xmin, ymin, xmax, ymax);
                Double xresolution = pcBroad.getValue().getBaseMapEnv().getWidth() / Math.pow(2, zBroad.getValue());
                Double yresolution = pcBroad.getValue().getBaseMapEnv().getHeight() / Math.pow(2, zBroad.getValue());
                RasterExtent re = MRasterizer.getRasterExtent(extent, xresolution, yresolution);

                File file = new File(outputPath);
                if (!file.exists())
                    file.mkdirs();

                List<Tuple2<String, String>> fs = new ArrayList<>();
                Map<OptionField, String> mfs = fieldValueInfoBroad.getValue();
                for (OptionField key : mfs.keySet()) {
                    if (key.getFieldType().equals("EMPTY")) continue;
                    fs.add(new Tuple2<>(key.getFieldName(), key.getFieldType()));
                }

                MultibandTile mbt = MRasterizer.rasterizeWithValue(gis, re, fs);
                String filepath = outputPath + File.separator + UUID.randomUUID().toString() + ".tiff";
                MRasterizer.write(mbt, extent, filepath);

                logger.info("Write tiff file: " + filepath);

                List<Double[]> result = new ArrayList<>();
                result.add(new Double[]{xmin, xmax, ymin, ymax});

                return result.iterator();
            }
        });

        Double[] extentsResult = extent.reduce(new Function2<Double[], Double[], Double[]>() {
            @Override
            public Double[] call(Double[] e1, Double[] e2) throws Exception {
                double xmin = e1[0] < e2[0] ? e1[0] : e2[0];
                double xmax = e1[1] < e2[1] ? e1[1] : e2[1];
                double ymin = e1[2] < e2[2] ? e1[2] : e2[2];
                double ymax = e1[3] < e2[3] ? e1[3] : e2[3];
                return new Double[]{xmin, ymin, xmax, ymax};
            }
        });

        logger.info(String.format("Tiff write success with file extent: {%.9f, %.9f, %.9f, %.9f}",
                extentsResult[0], extentsResult[1], extentsResult[2], extentsResult[3]));

        try {
            Class.forName(driver);
            conn = DriverManager.getConnection(url, user, password);
            stmt = conn.prepareStatement(sql);
            stmt.setString(1, publisher);
            stmt.setString(2, id);
            stmt.setString(3, layerName);
            stmt.setString(4, serviceURL);
            stmt.setString(5, keyWords);
            stmt.setString(6, "TMS");
            stmt.setString(7, topic);
            stmt.setString(8, "WGS84");
            stmt.setString(9, "经纬度");
            stmt.setString(10, String.format("%.9f,%.9f,%.9f,%.9f", extentsResult[0], extentsResult[1], extentsResult[2], extentsResult[3]));
            stmt.setTimestamp(11, Timestamp.valueOf(DateUtils.curDateTimeStrDefault()));
            stmt.setString(12, "");
            stmt.setString(13, "***/topography.png");
            stmt.setString(14, mArgs.getUserId());
            stmt.setString(15, field_info.toString());
            stmt.setString(16, "raster");
            stmt.setString(17, indexName);
            stmt.setString(18, "loading");
            stmt.executeUpdate();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } catch (SQLException e) {
            e.printStackTrace();
        }

        try {
            MIngest.run(jsc.sc(), outputPath, output, layerName, jsc.defaultParallelism());

            try {
                Class.forName(driver);
                conn = DriverManager.getConnection(url, user, password);

                int index = 0;
                for (OptionField key : fieldValueInfo.keySet()) {
                    if (key.getFieldType().equals("EMPTY")) continue;
                    stmt = conn.prepareStatement(sqlfield);
                    stmt.setString(1, UUID.randomUUID().toString());
                    stmt.setString(2, id);
                    stmt.setString(3, key.getFieldName());
                    stmt.setString(4, String.valueOf(index));
                    stmt.setString(5, key.getFieldType());
                    String fieldValue;
                    if (key.getFieldType().equals("DString")) {
                        Map<String, Integer> fValueTrans = valueTrans.get(key.getFieldName());
                        List<String> o = new ArrayList<>();
                        for (String k : fValueTrans.keySet()) {
                            Integer t = fValueTrans.get(k);
                            o.add(String.format("%s,%s", k, t));
                        }
                        fieldValue = String.join("##", o);
                    } else {
                        fieldValue = fieldValueInfo.get(key);
                    }
                    stmt.setString(6, fieldValue);

                    GeotrellisTileOptions fieldsColorMap = new GeotrellisTileOptions();
                    fieldsColorMap.setFieldType(key.getFieldType());
                    if (fieldsColorMap.getFieldType().equals("CDouble")) {
                        String[] vs = fieldValue.split("##");
                        int min = Double.valueOf(vs[0]).intValue();
                        int max = Double.valueOf(vs[1]).intValue() + 1;
                        fieldsColorMap.setMin(min);
                        fieldsColorMap.setMax(max);
                    } else {
                        ColorMap.Options options = new ColorMap.Options(edu.zju.gis.dbfg.model.scala.GeoTrellisTile.getClassBoundaryType(), 0x00000000, 0x00000000, false);
                        String[] vs = fieldValue.split("##");
                        if (key.getFieldType().equals("DString")) {
                            List<Double> temp = new ArrayList<>();
                            for (Double i = 0.0; i < vs.length; i++)
                                temp.add(i);
                            Map<String, String> fieldColor = new HashMap<>();
                            ColorMap tempCm = edu.zju.gis.dbfg.model.scala.GeoTrellisTile.getDiscreteRandomColorMap(temp, options);
                            String[] pairs = tempCm.breaksString().split(";");
                            for (int i = 1; i < pairs.length + 1; i++) {
                                String name = "N";
                                for (String p : vs) {
                                    if (p.split(",")[1].equals(String.valueOf(i))) {
                                        name = p.split(",")[0];
                                        break;
                                    }
                                }
                                fieldColor.put(name, pairs[i - 1].split(":")[1]);
                            }
                            fieldsColorMap.setColorMap(fieldColor);
                        } else {
                            List<Double> catchValue = new ArrayList<>();
                            for (String v : vs) {
                                catchValue.add(Double.valueOf(v));
                            }
                            Map<String, String> fieldColor = new HashMap<>();
                            ColorMap tempCm = edu.zju.gis.dbfg.model.scala.GeoTrellisTile.getDiscreteRandomColorMap(catchValue, options);
                            String[] pairs = tempCm.breaksString().split(";");
                            for (int i = 0; i < pairs.length; i++) {
                                fieldColor.put(pairs[i].split(":")[0], pairs[i].split(":")[1]);
                            }
                            fieldsColorMap.setColorMap(fieldColor);
                        }
                    }
                    fieldsColorMap.setColorRampIndex(1);
                    ColorMap.Options options = new ColorMap.Options(edu.zju.gis.dbfg.model.scala.GeoTrellisTile.getClassBoundaryType(), 0x00000000, 0x00000000, false);
                    ColorMap cm = null;
                    switch (fieldsColorMap.getFieldType()) {
                        case "CDouble":
                            BigInteger sc = new BigInteger(fieldsColorMap.getStartColor(), 16);
                            BigInteger ec = new BigInteger(fieldsColorMap.getEndColor(), 16);
                            int startColor = sc.intValue();
                            int endColor = ec.intValue();
                            int min = fieldsColorMap.getMin();
                            int max = fieldsColorMap.getMax();
                            int stops = fieldsColorMap.getStops();
                            double gap = (max - min) * 1.0 / stops;
                            double[] stopVals = new double[stops];
                            for (int i = 0; i < stopVals.length; i++) {
                                stopVals[i] = (min + i * gap);
                            }
                            cm = edu.zju.gis.dbfg.model.scala.GeoTrellisTile.getCustomContinouesColorMap(startColor, endColor, min, max, stopVals, options);
                            break;
                        case "DString":
                            cm = catchColorMapStr(fieldValue, fieldsColorMap, options);
                            break;
                        case "DDouble":
                            cm = catchColorMap(fieldsColorMap, options);
                            break;
                    }
                    stmt.setString(7, cm.breaksString());
                    stmt.executeUpdate();
                    index++;
                }

                stmt = conn.prepareStatement(updateSql);
                stmt.setString(1, "succeeded");
                stmt.setString(2, id);
                stmt.executeUpdate();
            } catch (ClassNotFoundException e) {
                e.printStackTrace();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        } catch (Exception ex) {
            System.out.println(ex.getMessage());
            try{
                Class.forName(driver);
                conn = DriverManager.getConnection(url, user, password);
                stmt = conn.prepareStatement(updateSql);
                stmt.setString(1, "failed");
                stmt.setString(2, id);
                stmt.executeUpdate();
            }catch (ClassNotFoundException e) {
                e.printStackTrace();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    private static ColorMap catchColorMapStr(String fieldValue, GeotrellisTileOptions fieldsColorMap, ColorMap.Options options) {
        Map<Double, Integer> colorMap = new HashMap<>();
        Set<String> keys = fieldsColorMap.getColorMap().keySet();

        String[] items = fieldValue.split("##");
        Map<String, String> pairs = new HashMap<>();
        for (String item : items) {
            String[] kv = item.split(",");
            pairs.put(kv[0], kv[1]);
        }

        for (String key : keys) {
            String k = pairs.get(key);
            String v = fieldsColorMap.getColorMap().get(key);
            BigInteger bi = new BigInteger(v, 16);
            int temp = bi.intValue();
            colorMap.put(Double.valueOf(k), temp);
        }
        return edu.zju.gis.dbfg.model.scala.GeoTrellisTile.transformColorMap(colorMap, options);
    }

    private static ColorMap catchColorMap(GeotrellisTileOptions fieldsColorMap, ColorMap.Options options) {
        Map<Double, Integer> colorMap = new HashMap<>();
        Set<String> keys = fieldsColorMap.getColorMap().keySet();
        while (keys.iterator().hasNext()) {
            String val = keys.iterator().next();
            colorMap.put(Double.valueOf(val), Integer.valueOf(fieldsColorMap.getColorMap().get(val)));
        }
        return edu.zju.gis.dbfg.model.scala.GeoTrellisTile.transformColorMap(colorMap, options);
    }
}
